package com.woniu.websocket;

import org.springframework.stereotype.Component;

import javax.websocket.OnClose;
import javax.websocket.OnError;
import javax.websocket.OnOpen;
import javax.websocket.Session;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;

@Component
@ServerEndpoint("/ws/notfication/{role}")
public class NotificationWebSocket {
    // 按角色存储会话
    private static final ConcurrentHashMap<String, ConcurrentHashMap<Session, String>> roleRooms = new ConcurrentHashMap<>();

    @OnOpen
    public void onOpen(Session session, @PathParam("role") String role) {
        roleRooms.computeIfAbsent(role, k -> new ConcurrentHashMap<>()).put(session, "");
    }

    @OnClose
    public void onClose(Session session, @PathParam("role") String role) {
        ConcurrentHashMap<Session, String> sessions = roleRooms.get(role);
        if (sessions != null) {
            sessions.remove(session);
        }
    }

    @OnError
    public void onError(Session session, Throwable e) {
        e.printStackTrace();
    }

    /**
     * 发送消息给指定角色的所有在线用户
     */
    public static void sendToRole(String role, String message) {
        ConcurrentHashMap<Session, String> sessions = roleRooms.get(role);
        if (sessions != null) {
            sessions.keySet().forEach(s -> {
                if (s.isOpen()) {
                    try {
                        s.getBasicRemote().sendText(message);
                    } catch (IOException ignored) {}
                }
            });
        }
    }
}
